# File: util/trace_validator.py

import numpy as np
from deadlock_env import STATE_BLOCKED

# --- [新增] 定义更多奖励/惩罚常量 ---
REWARD_CORRECT_WAIT = 0.05  # 提高过程奖励
PENALTY_WRONG_WAIT = -0.05
REWARD_CORRECT_HOLD = 0.05  # 提高过程奖励
PENALTY_WRONG_HOLD = -0.05
REWARD_VALID_CYCLE = 0.2
PENALTY_INVALID_CYCLE = -0.2
PENALTY_MALFORMED_CYCLE = -0.25

# --- [核心] 新增决策相关的奖励 ---
# 奖励1: 模型成功生成了一个语法上可解析的决策
REWARD_VALID_DECISION = 0.3
# 奖励2: 模型选择的牺牲品进程确实是死锁的一部分 (一个更明智的决策)
REWARD_CORRECT_VICTIM = 0.5


class TraceValidator:
    """
    分析模型生成的思维链，并与真实环境状态对比，以提供密集的奖励信号。
    """

    def __init__(self, symbol_manager):
        if not hasattr(symbol_manager, 'parse_trace_to_assertions'):
            raise AttributeError("The provided 'symbol_manager' instance must have a "
                                 "'parse_trace_to_assertions' method.")
        self.sm = symbol_manager

    def _build_wait_for_graph(self, current_obs: dict) -> dict:
        """从当前观测中构建一个 P->P 的等待图。"""
        num_processes = current_obs['allocation'].shape[0]
        wait_for_graph = {}
        resource_holder = {r_id: p_id for p_id in range(num_processes) for r_id, count in enumerate(current_obs['allocation'][p_id]) if count > 0}

        for p_wait in range(num_processes):
            if current_obs['processes_status'][p_wait] == STATE_BLOCKED:
                requested_r_indices = np.where(current_obs['need'][p_wait] > 0)[0]
                if len(requested_r_indices) > 0:
                    r_needed = requested_r_indices[0]
                    if r_needed in resource_holder:
                        p_holds = resource_holder[r_needed]
                        wait_for_graph[p_wait] = p_holds
        return wait_for_graph

    def validate_and_reward(self, trace_tokens: np.ndarray, current_obs: dict) -> float:
        """
        验证整个思维链并计算总奖励值。
        """
        reward = 0.0
        parsed_assertions = self.sm.parse_trace_to_assertions(trace_tokens)
        wait_for_graph = self._build_wait_for_graph(current_obs)

        # 找出所有真正参与到阻塞关系中的进程ID
        involved_processes = set(wait_for_graph.keys()) | set(wait_for_graph.values())

        for assertion in parsed_assertions:
            # --- 1. 验证 "WAIT" 关系 ---
            if assertion['type'] == 'WAIT':
                p_id, r_id = assertion['process'], assertion['resource']
                is_blocked = (current_obs['processes_status'][p_id] == STATE_BLOCKED)
                is_waiting_for_r = (current_obs['need'][p_id, r_id] > 0)
                if is_blocked and is_waiting_for_r:
                    reward += REWARD_CORRECT_WAIT
                else:
                    reward += PENALTY_WRONG_WAIT

            # --- 2. 验证 "HELD_BY" 关系 ---
            elif assertion['type'] == 'HELD_BY':
                r_id, p_id = assertion['resource'], assertion['process']
                if current_obs['allocation'][p_id, r_id] > 0:
                    reward += REWARD_CORRECT_HOLD
                else:
                    reward += PENALTY_WRONG_HOLD

            # --- 3. 验证 "CYCLE" ---
            elif assertion['type'] == 'CYCLE':
                path = assertion['path']
                if not isinstance(path, list) or len(path) < 2:
                    reward += PENALTY_MALFORMED_CYCLE
                    continue

                is_valid_cycle = True
                for i in range(len(path)):
                    p_wait = path[i]
                    p_hold = path[(i + 1) % len(path)]
                    if wait_for_graph.get(p_wait) != p_hold or current_obs['processes_status'][p_wait] != STATE_BLOCKED:
                        is_valid_cycle = False
                        break

                if is_valid_cycle:
                    reward += REWARD_VALID_CYCLE
                else:
                    reward += PENALTY_INVALID_CYCLE

            # --- 4. [核心新增] 验证 "DECIDE" ---
            elif assertion['type'] == 'DECIDE':
                # 基础奖励：只要能成功解析出一个决策，就给予奖励
                reward += REWARD_VALID_DECISION

                victim_pid = assertion['process']
                # 额外奖励：如果选择的牺牲品是真正参与死锁的进程，给予额外奖励
                if victim_pid in involved_processes:
                    reward += REWARD_CORRECT_VICTIM

        return reward
